namespace NBody.Cosmology
{
    using System;
    using System.Reflection;
    
    public delegate double TransferFunction(double k);
    
    public enum Normalization { Sigma8, CMB, None };
    
    // everything here is given at z = 0.  At other redshifts, scale.
    
    public class Spectrum
    {
        private Model model;
        private double n = 1;
        private double peak = 4;
        private double pow_amp = 1.0;
        private double sigma_8 = 0.63;
        private double smooth_scale = 0.0;
        private Normalization normalization;
        private TransferFunction transfer_function;
        
        private double Rstar, gamma, theta, xi_0;
        
        public Spectrum(Model model, TransferFunction f)
        {
            this.model = model;
            model.ModelChange += OnModelChange;
            this.transfer_function = f;
            normalization = Normalization.None;
        }
        
        private void OnModelChange(object o, EventArgs args)
        {
            CalcAmp();
        }
        
        public double Index
        {
            get { return n; }
            set { n = value; CalcAmp(); }
        }
        
        public Model Model
        {
            get { return model; }
        }
        
        public double Sigma8
        {
            get { return sigma_8; }
            set { sigma_8 = value;  CalcAmp();}
        }
        
        public double SmoothingScale
        {
            get { return smooth_scale; }
            set { smooth_scale = value;  CalcAmp();}
        }
        
        public double PeakHeight
        {
            get { return peak; }
            set { peak = value; }
        }
        
        public double Amplitude
        {
            get { return pow_amp; }
            set { if (normalization == Normalization.None) { pow_amp = value; CalcAmp(); } }
        }
        
        public Normalization Normalization
        {
            get { return normalization; }
            set { normalization = value; CalcAmp(); }
        }
        
        public TransferFunction TransferFunction
        {
            get { return transfer_function; }
        }
        
        public double PowerSpectrum(double k)
        {
            // switch to mpc-1
            double q = 1000.0 * k / model.HParam / model.HParam / model.OmegaM;
            double t = transfer_function(q);
            double smooth = 1.0;
            if (smooth_scale > 0.0)
                smooth = Math.Exp(-0.5 * k * k * smooth_scale * smooth_scale);
            double p = pow_amp * Math.Pow(k, n) * t * t * smooth;
            return p;
        }
        
        public double CorrelationFunction(double r)
        {
            Gsl.Integration i = new Gsl.Integration(1000);
            i.Function = new Gsl.Function(CorrFuncInt);
            i.Parameters = r;
            i.ExecuteQAGIU(0.0);
            
            double xi = 1.0 / 2.0 / Math.PI / Math.PI / r * i.Result;
            return xi;
        }
        
        private double CorrFuncInt(double k, object radius)
        {
            double r = (double)radius;
            double p = PowerSpectrum(k);
            double i = p * k * Math.Sin(k * r);
            return i;
        }
        
        private double SpectralMoment(int l)
        {
            Gsl.Integration i = new Gsl.Integration(1000);
            i.Function = new Gsl.Function(SpectralMomentInt);
            i.Parameters = l;
            i.ExecuteQAGIU(0.0);
            
            double ans = 1.0 / 2.0 / Math.PI / Math.PI * i.Result;
            return Math.Sqrt(ans);
        }
        
        private double SpectralMomentInt(double k, object o)
        {
            double l = (int)o;
            return PowerSpectrum(k) * Math.Pow(k, 2.0 * (l+1.0));
        }
        
        private double CorrFuncGrad2(double r)
        {
            Gsl.Integration i = new Gsl.Integration(1000);
            i.Function = new Gsl.Function(CorrFuncGrad2Int);
            i.Parameters = r;
            i.ExecuteQAGIU(0.0);
            
            double ans = -1.0 / (6.0 * Math.PI * Math.PI * r) * i.Result;
            return ans;
        }
        
        private double CorrFuncGrad2Int(double k, object o)
        {
            double r = (double)o;
            double p = PowerSpectrum(k);
            double y = k*k* Math.Sin(k * r);
            double i = p * k * y;
            return i;
        }
        
        public double Delta(double x)
        {
            double r = x;
            double xi = CorrelationFunction(r);
            double xi_diff2 = CorrFuncGrad2(r);
            double d = peak * xi / Math.Sqrt(xi_0) - theta / gamma / (1.0 - gamma * gamma) * (gamma*gamma * xi + Rstar * Rstar * xi_diff2) / Math.Sqrt(xi_0);
            return d;
        }
        
        // simplest transfer function
        public static double TransferSimple(double k)
        {
            return 1.0;
        }
        
        // BBKS SCDM transfer function
        // q should be in Mpc h^2
        public static double TransferBBKS(double q)
        {
            double a0 = 2.34 * q;
            double a1 = 3.89 * q;
            double a2 = 16.1 * q;
            double a3 = 5.46 * q;
            double a4 = 6.71 * q;
            double T = Math.Log(1.0 + a0) / a0 / Math.Sqrt(Math.Sqrt(1.0 + a1 + a2*a2 + a3*a3*a3 + a4*a4*a4*a4));
            
            return T;
        }
        
        // Lambda CDM transfer function from Dave's thesis
        public static double TransferLambdaCDM(double q)
        {
            double b1 = 1.04e4;
            double b2 = 1.96e5;
            
            return 1.0 / (1.0 + b1 * q * q + b2 * Math.Pow(q, 3.5));
        }
        
        public static Spectrum SCDM
        {
            get
            {
                return new Spectrum(Model.SCDM, new TransferFunction(TransferBBKS));
            }
        }
        
        public static Spectrum FromString(string cosmo_string)
        {
            Spectrum cosmo;
            try
            {
                Type cosmo_type = typeof(Spectrum);
                cosmo = (Spectrum)cosmo_type.InvokeMember(cosmo_string, BindingFlags.Public | BindingFlags.Static | BindingFlags.GetProperty, null, null, null);
            }
            catch (Exception e)
            {
                Console.Error.WriteLine("Can't find cosmology spectrum for \"{0}\"; falling back to good old SCDM", cosmo_string);
                Console.Error.WriteLine("  Problem: {0}", e.Message);
                cosmo = Spectrum.SCDM;
            }
            
            return cosmo;
        }
        
        //
        // private functions
        //
        
        private void CalcAmp()
        {
            switch (normalization)
            {
                case Normalization.None:
                    break;
                case Normalization.CMB:
                    Console.Error.WriteLine("CMB normalization not supported right now.");
                    break;
                case Normalization.Sigma8:
                    pow_amp = 1.0;
                    pow_amp = CalcPowAmpSigma8();
                    break;
            }
            
            Console.Error.WriteLine("Power spectrum normalized with {0} to an amplitude {1}", normalization, pow_amp);
            
            // now do spectra moments for delta calc ... but only if smooth_scale > 0.0, since we
            // won't need them otherwise
            if (smooth_scale <= 0.0) return;
            
            double sigma0 = SpectralMoment(0);
            double sigma1 = SpectralMoment(1);
            double sigma2 = SpectralMoment(2);
            Rstar = Math.Sqrt(3.0) * sigma1 / sigma2;
            gamma = sigma1*sigma1 / sigma2 / sigma0;
            Console.Error.WriteLine("Smoothing power spectrum over {0} kpc", smooth_scale);
            Console.Error.WriteLine("  Rstar = {0}", Rstar);
            Console.Error.WriteLine("  gamma = {0}", gamma);
            
            // and a couple other things
            xi_0 = CorrelationFunction(1e-5);
            theta = (3.0 * (1.0 - gamma * gamma) + (1.216 - 0.9 * Math.Pow(gamma, 4.0)) * Math.Exp(-0.5 * gamma * Math.Pow(0.5 * peak * gamma, 2.0))) / (Math.Sqrt(3.0 * (1.0 - gamma * gamma) + 0.45 + Math.Pow(peak * gamma / 2.0, 2.0)) + 0.5 * peak * gamma);
        }
        
        private double CalcPowAmpSigma8()
        {
            double R = 8000.0 / model.HParam; // 8 mpc/h 
            
            Gsl.Integration i = new Gsl.Integration(1000);
            i.Function = Sigma2Int;
            i.Parameters = R;
            
            i.ExecuteQAGIU(0.0);
            double sigma2 = 1.0 / 2.0 / Math.PI / Math.PI * i.Result;
            //double sigma2 = 4.0 * Math.PI * i.Result;
            return sigma_8 * sigma_8 / sigma2;
        }
        
        private double Sigma2Int(double k, object par)
        {
            double x = k * (double)par;
            double w = 3.0 * (Math.Sin(x) - x * Math.Cos(x)) / (x * x * x);
            double p = PowerSpectrum(k);
            return p * w * w * k * k;
        }
    }
}
